import html
import threading
import time

from flask import request, jsonify, Response

from Common import *
from ShieldNet import ModelTrainer, ModelPlotSaver, DDoSType, ModelCreatorConfig, DataLoader, \
    stream_log_handler, PCAPProcessor

from ShieldNet.Detector import *

model_type: str = DDoSType.DrDoS_DNS.value


@app.route("/DetectApi/Model/LoadModel")
def load_model():
    global model_type
    model_type = DDoSType(request.get_json(silent=True).get("model_type", DDoSType.DrDoS_DNS.value))
    return jsonify({"status": "success", "message": f"模型类型已设置为{model_type.name}"}), 200

@app.route("/DetectApi/Model/GetModelType")
def get_model_type():
    return jsonify({"status": "success", "model_type": model_type}), 200

@app.route("/DetectApi/Model/SetPlotSavePath")
def set_plot_save_path():
    set_plot_save_path(request.get_json(silent=True).get("plot_save_path", plot_save_path))

@app.route("/DetectApi/Model/TrainModel", methods=["POST"])
def train_model():
    json_data = request.get_json(silent=True)
    if json_data is None or "epochs" not in json_data or "train_data_path" not in json_data:
        return jsonify({"status": "error", "message": "错误的请求格式，必须带有训练次数[epochs][train_data_path]参数"}), 400
    else:
        # 训练模型
        train_data_path = json_data.pop("train_data_path")
        is_clear_pth_and_pkl = json_data.pop("is_clear_pth_and_pkl", True)
        train = lambda: (ModelTrainer(
            model_plot_saver=ModelPlotSaver(
                plot_train_save_path=os.path.join(plot_save_path, model_type, str(json_data.pop("train_png_name", "train.png"))),
                plot_cm_save_path=os.path.join(plot_save_path, model_type, str(json_data.pop("cm_png_name", "cm.png"))),
                plot_pr_save_path=os.path.join(plot_save_path, model_type, str(json_data.pop("pr_png_name", "pr.png"))),
            ),
            model_creator_config=ModelCreatorConfig(
                model_path=str(json_data.pop("model_path", ".\\model")),
                model_name=str(json_data.pop("model_name", "Model")),
                suffix="mod",
                pth_path=os.path.join(common_save_path, model_type, str(json_data.pop("pth_name", "model.pth"))),
                static_pth_path=os.path.join(common_save_path, model_type, str(json_data.pop("static_pth_name", "static_model.pth"))),
                pkl_path=os.path.join(common_save_path, model_type, str(json_data.pop("pkl_name", "model.pkl"))),
            ),
            classify_name="Label",
            classify_loader=["BENIGN", model_type],
            data_loader=DataLoader(chunk_size=int(json_data.pop("chunk_size", 100000))),
            detector=eval(json_data.pop("detector", "DDoSDetector")),
            **json_data
        ).load_data(train_data_path)
         .deal_with()
         .divide_dataset()
         .init_model(**json_data.pop("init_model_kwargs", {}))
         .train()
         .show_train_plot()
         .model_evaluate()
         .plot_pr_curve()
         .clear_model(is_clear_pth_and_pkl))
        threading.Thread(target=train).start()
        return jsonify({"status": "success", "message": "模型训练已经启动，请查看日志"}), 200


@app.route("/DetectApi/Model/TrainerLog", methods=["POST"])
def trainer_log():
    def generate():
        sent_logs_count = 0
        while True:
            current_logs = stream_log_handler.get_logs()

            # 只发送新日志
            if len(current_logs) > sent_logs_count:
                new_logs = current_logs[sent_logs_count:]
                for log in new_logs:
                    # HTML转义防止XSS，同时保留格式
                    escaped_log = html.escape(log)
                    yield f"data: {escaped_log}\n\n"
                sent_logs_count = len(current_logs)

            time.sleep(0.1)

    return Response(
        generate(),
        mimetype="text/event-stream",
        headers={
            'Cache-Control': 'no-cache',
            'Connection': 'keep-alive',
            'X-Accel-Buffering': 'no'  # 禁用Nginx缓冲
        })

@app.route("/DetectApi/Model/run", methods=["POST"])
def run():
    model_path = request.get_json(silent=True).get("model_path", "")
    pcap_file = request.get_json(silent=True).get("pcap_file", "")
    window_size = request.get_json(silent=True).get("window_size", 30)
    attack_threshold = request.get_json(silent=True).get("attack_threshold", 0.7)
    if not os.path.exists(model_path) or not os.path.exists(pcap_file):
        return jsonify({"status": "error", "message": "模型或PCAP文件不存在"}), 400
    try:
        processor = PCAPProcessor(
            model_path=model_path,
            window_size=window_size,
            output_dir="../detection_results",
            attack_threshold=attack_threshold
        )
        results = processor.process_pcap(pcap_file)
        return jsonify({"status": "success", "results": results}), 200
    except Exception as e:
        logger.error(f"处理失败: {type(e).__name__}: {str(e)}")
        return jsonify({"status": "error", "message": f"处理失败: {type(e).__name__}: {str(e)}"}), 500